/*
* @Author: ystlong
* @Date:   2018-06-29 18:55:11
* @Last Modified by:   ystlong
* @Last Modified time: 2018-06-30 03:02:59
*/

#include <stdio.h>
#include <stdlib.h>
#include <elf.h>
#include <assert.h>
#include <string.h>
#include "parse_elf.h"

#ifdef COMPILE_ELF64

// #define Elf_Byte        Elf_Byte    
#define ElfN_Addr       Elf64_Addr        
#define ElfN_Off        Elf64_Off       
#define ElfN_Section    Elf64_Section           
// #define ElfN_Versym     Elf64_Versym          
// #define ElfN_Half       Elf64_Half        
// #define ElfN_Sword      Elf64_Sword         
// #define ElfN_Word       Elf64_Word        
// #define ElfN_Sxword     Elf64_Sxword          
// #define ElfN_Xword      Elf64_Xword         
#define ElfN_Shdr       Elf64_Shdr
#define ElfN_Ehdr       Elf64_Ehdr
#define ElfN_Sym        Elf64_Sym
#define ElfN_parse		elf64_parse

#else 
#ifdef COMPILE_ELF32

// #define Elf_Byte        Elf_Byte    
#define ElfN_Addr       Elf32_Addr        
#define ElfN_Off        Elf32_Off       
#define ElfN_Section    Elf32_Section           
// #define ElfN_Versym     Elf32_Versym          
// #define ElfN_Half       Elf32_Half        
// #define ElfN_Sword      Elf32_Sword         
// #define ElfN_Word       Elf32_Word        
// #define ElfN_Sxword     Elf32_Sxword          
// #define ElfN_Xword      Elf32_Xword         
#define ElfN_Shdr       Elf32_Shdr
#define ElfN_Ehdr       Elf32_Ehdr
#define ElfN_Sym        Elf32_Sym
#define ElfN_parse		elf32_parse

#else
#error "build mustbe  define COMPILE_ELF32 or COMPILE_ELF64"
#endif // COMPILE_ELF32
#endif //COMPILE_ELF64

typedef struct 
{
	ElfN_Shdr section;
	unsigned char *data;
}SectionData;


// static int sec

static SectionData *new_elf_sectiondata()
{
	SectionData * new_sec_data = (SectionData*)malloc(sizeof(SectionData));
	new_sec_data->data = NULL;
	return new_sec_data;
}

static void del_elf_sectiondata(SectionData *sec_data)
{
	if(sec_data != NULL) {
		if(sec_data->data != NULL)
			free(sec_data->data);
		free(sec_data);
	} 
}

static int get_section_et_by_inx(FILE *elf_fs, ElfN_Ehdr *elf_header, uint16_t section_inx, ElfN_Shdr * elf_shdr)
{
	ElfN_Off section_off = elf_header->e_shoff + elf_header->e_shentsize*section_inx;
	fseek(elf_fs, section_off, SEEK_SET);
	fread(elf_shdr, sizeof(*elf_shdr), 1, elf_fs);
	return 0;
}

static int read_section_et_data(FILE *elf_fs, SectionData *section_data)
{
	section_data->data = (unsigned char*)malloc(section_data->section.sh_size);
	if (section_data->data == NULL)
		return 1;

	fseek(elf_fs, section_data->section.sh_offset, SEEK_SET);
	fread(section_data->data, section_data->section.sh_size, 1, elf_fs);
	return 0;
}

static int read_section(FILE *elf_fs, ElfN_Ehdr *elf_header, 
						uint16_t section_inx, SectionData * section_data)
{
	get_section_et_by_inx(elf_fs, elf_header, section_inx, &section_data->section);
	// read section data
	read_section_et_data(elf_fs, section_data);
	return 0;
}

static int read_section_strtab(FILE *elf_fs, ElfN_Ehdr *elf_header, SectionData *section_data)
{
	return read_section(elf_fs, elf_header, elf_header->e_shstrndx, section_data);
}

static char * get_section_name(ElfN_Shdr *elf_shdr, SectionData *sec_startab_section)
{
	return (char *)sec_startab_section->data + elf_shdr->sh_name;
}

static int find_section_by_name(FILE *elf_fs, ElfN_Ehdr *elf_header, 
	SectionData *sec_startab_section, SectionData *dest_section,
	char * sec_name)
{
	ElfN_Shdr *elf_shdr = &dest_section->section;
	for(int i=0; i<elf_header->e_shnum; i++) {
		get_section_et_by_inx(elf_fs, elf_header, i, elf_shdr);
		char *name = get_section_name(elf_shdr, sec_startab_section);
		if (strcmp(name, sec_name) == 0) {
			return 0;
		}
	}
	return 1;
}

static int get_strtab_section(FILE *elf_fs, ElfN_Ehdr *elf_header, 
								SectionData *sec_startab_section, SectionData *dest_section)
{
	int ret = find_section_by_name(elf_fs, elf_header, sec_startab_section, dest_section, ".strtab");
	if (!ret) {
		return read_section_et_data(elf_fs, dest_section);
	}
	return ret;
}

static void info_section(SectionData *sec_data, SectionData *sec_startab_section)
{
	if(sec_startab_section)
		printf("section %s, ", sec_startab_section->data + sec_data->section.sh_name);
	// printf("sh_name, %d, ", sec_data->section.sh_name);
	// printf("sh_type, %d, ", sec_data->section.sh_type);
	// printf("sh_flags, %d, ", sec_data->section.sh_flags);
	printf("sh_addr, 0x%x, ", sec_data->section.sh_addr);
	printf("sh_offset, %d, ", sec_data->section.sh_offset);
	printf("sh_size, %d, ", sec_data->section.sh_size);
	// printf("sh_link, %d, ", sec_data->section.sh_link);
	// printf("sh_info, %d, ", sec_data->section.sh_info);
	// printf("sh_addralign, %d, ", sec_data->section.sh_addralign);
	printf("sh_entsize, %d, ", sec_data->section.sh_entsize);
	printf("\n");
}

static void info_symtab_et(ElfN_Sym *sym_et, SectionData *strtab_section)
{
	if(strtab_section != NULL)
		printf("st name: %s, ", strtab_section->data+sym_et->st_name);
	printf("st_name: %d, ", sym_et->st_name);
	printf("st_info: %d, ", sym_et->st_info);
	printf("st_other: %d, ", sym_et->st_other);
	printf("st_shndx: %d, ", sym_et->st_shndx);
	printf("st_value: 0x%x, ", sym_et->st_value);
	printf("st_size: %d, ", sym_et->st_size);
	printf("\n");
}

static SectionData *get_section_by_name(FILE *elf_fs, ElfN_Ehdr *elf_header,
										SectionData *sec_startab_section, char *name)
{
	SectionData *tmp_sec_data = new_elf_sectiondata();

	int ret = find_section_by_name(elf_fs, elf_header, sec_startab_section, tmp_sec_data, name);
	if (!ret) {
		ret = read_section_et_data(elf_fs, tmp_sec_data);
	}
	if (ret) {
		del_elf_sectiondata(tmp_sec_data);
		// printf("warn: %s symtab not found\n", name);
		// exit(1);
		return NULL;
	}else{
		// info_section(strtab_section, sec_startab_section);
		// int psize = 0;
		// for(int j=0; j<strtab_section->section.sh_size; j+= psize) {
		// 	printf("---: pos: %d, ", j);
		// 	psize = printf("%s\n", strtab_section->data + j);
		// }
	}
	return tmp_sec_data;
}

static char *get_strtab_et(SectionData *str_section, uint32_t inx)
{
	return (char *)str_section->data + inx;
}

static ElfN_Sym *find_symtab_et_by_name(SectionData *symtab_section, SectionData *strtab_section, char *name)
{
	ElfN_Sym *tmp_sym = (ElfN_Sym *)(symtab_section->data);
	assert(sizeof(ElfN_Sym) == symtab_section->section.sh_entsize);

	for(int j=0; j<symtab_section->section.sh_size/symtab_section->section.sh_entsize; j++) {
		// if ((tmp_sym+j)->st_info == DT_RELSZ)
		switch(tmp_sym->st_info){
			// case DT_PLTGOT:
			case DT_PLTRELSZ:
			case DT_RELA:
			case DT_RELASZ:
			case DT_RELAENT:
			case DT_REL:
			case DT_RELSZ:
			case DT_RELENT:
			case DT_PLTREL:
			case DT_TEXTREL:
			case DT_JMPREL:
			// case DT_RELACOUNT:
			// case DT_RELCOUNT:
				// char *name = get_strtab_et(strtab_section, tmp_sym->st_name);
				if (strcmp(name, get_strtab_et(strtab_section, tmp_sym->st_name)) == 0){
					return tmp_sym;
				}
			// default:
		}
		tmp_sym++;
	}
	return NULL;
}

static ElfN_Off addr_2_offset(FILE *elf_fs, ElfN_Ehdr *elf_header, ElfN_Addr addr)
{
	SectionData *tmp_sec_data = new_elf_sectiondata();
	ElfN_Shdr *elf_shdr = &tmp_sec_data->section;
	ElfN_Off file_off = -1;
	for(int i=0; i<elf_header->e_shnum; i++) {
		if(get_section_et_by_inx(elf_fs, elf_header, i, elf_shdr)){
			printf("error get_section_et_by_inx");
			exit(3);
		}
		// info_section(elf_shdr, NULL);
		if(addr >= elf_shdr->sh_addr && addr < (elf_shdr->sh_addr+elf_shdr->sh_size)){
			// find addr
			ElfN_Off rel_off = addr - elf_shdr->sh_addr;
			file_off = elf_shdr->sh_offset + rel_off;
		}
	}
	del_elf_sectiondata(tmp_sec_data);
	return file_off;
}

static void print_strtab_all_data(SectionData *str_section)
{
	int psize = 0;
	for(int j=0; j<str_section->section.sh_size; j+= psize) {
		printf("---: pos: %d, ", j);
		psize = printf("%s\n", str_section->data + j);
	}
}

static long _elf_parse(FILE *elf_fs, int patch_type, char *target, char *strtab_name, char *symtab_name)
{
	ElfN_Ehdr elf_header;
	rewind(elf_fs);
	fread(&elf_header, sizeof(elf_header), 1, elf_fs);

	SectionData *sec_startab_section = new_elf_sectiondata();
	if(read_section_strtab(elf_fs, &elf_header, sec_startab_section)){
		printf("read strtab section error\n");
		return -1;
		// exit(1);
	}
	// print_strtab_all_data(sec_startab_section);

	// 全局字符串， 包括： 函数名， 文件名...
	// SectionData *strtab_section = get_section_by_name(elf_fs, &elf_header, sec_startab_section, ".strtab");
	// SectionData *strtab_section = get_section_by_name(elf_fs, &elf_header, sec_startab_section, ".dynstr");
	SectionData *strtab_section = get_section_by_name(elf_fs, &elf_header, sec_startab_section, strtab_name);
	if (strtab_section == NULL) {
		return -2;
	}
	// 符号表， 包含函数对应地址
	// SectionData *symtab_section = get_section_by_name(elf_fs, &elf_header, sec_startab_section, ".symtab");
	// SectionData *symtab_section = get_section_by_name(elf_fs, &elf_header, sec_startab_section, ".dynsym");
	SectionData *symtab_section = get_section_by_name(elf_fs, &elf_header, sec_startab_section, symtab_name);
	if(symtab_section == NULL) {
		return -3;
	}

	// info_section(symtab_section);
	ElfN_Addr tag_addr = -1;

	switch(patch_type){
		case PATCH_FUN_NAME:
		{
			ElfN_Sym *main_sym = find_symtab_et_by_name(symtab_section, strtab_section, target);
			if (main_sym == NULL) {
				printf("error: fn %s not find in %s section\n", target, symtab_name);
				return -4;
				// exit(3);
			}else{
				// info_symtab_et(main_sym, strtab_section);
				tag_addr = main_sym->st_value;
				printf("find function %s at addr 0x%x\n", target, tag_addr);
			}
		}
		break;
		case PATCH_ADDR:
		{
			tag_addr = strtoul(target, NULL, 16);
			printf("input addr: 0x%x\n", tag_addr);
		}
	}

 	ElfN_Off file_off = addr_2_offset(elf_fs, &elf_header, tag_addr);
 	if (file_off == -1){
 		printf("error: addr 0x%x not find in and section\n", tag_addr);
 		return -6;
 		// exit(4);
 	}
 	printf("find addr 0x%x at offset %lu\n", tag_addr, file_off);

 	fseek(elf_fs, file_off, SEEK_SET);
 	for(int i=0; i<2; i++) {
	 	uint64_t data;
	 	fread(&data, sizeof(data), 1, elf_fs);
	 	for(int i=0; i<sizeof(data); i++)
	 		printf("%02x", ((unsigned char *)&data)[i]);
 	}
	printf("\n");

	del_elf_sectiondata(sec_startab_section);
	del_elf_sectiondata(strtab_section);
	del_elf_sectiondata(symtab_section);
	return file_off;
}

long ElfN_parse(FILE *elf_fs, int patch_type, char *target)
{
	long offset = 0;
	offset = _elf_parse(elf_fs, patch_type, target, ".strtab", ".symtab");
	if(offset < 0)
		offset = _elf_parse(elf_fs, patch_type, target, ".dynstr", ".dynsym");
	if(offset < 0) {
		printf("find target failed\n");
		exit(33);
	}
	return offset;
}
